import json
import math


alpha = 50.0

def _load_model(file):
    with open(file, 'r') as f:
        type_map = json.loads(f.read())
        trans_map = {}
        model = {}
        for k in type_map:
            ks = k.split(':')
            if len(ks) > 1:
                first_action = ks[0]
                rest_action = ':'.join(ks[1:])
                if first_action in trans_map:
                    trans_map[first_action][rest_action] = type_map[k]
                else:
                    trans_map[first_action] = {rest_action: type_map[k]}
        for k in trans_map:
            value_sum = float(sum(trans_map[k].values())) + 25.0 * alpha
            for _k in trans_map[k]:
                model[':'.join([_k, k])] = math.log((trans_map[k][_k] + alpha) / value_sum)
            model[':'.join([k, 'default'])] = math.log(alpha / value_sum)
        return model


def load_model():
    return {
        'type2': _load_model('./type_map2.json'),
        'type3': _load_model('./type_map3.json'),
    }


if __name__ == '__main__':
    print(_load_model('ttype_map3.json'))
